import socket
import asyncio
from asyncio import AbstractEventLoop
from typing import List
import signal


async def echo(connection: socket, loop: AbstractEventLoop) -> None:
    try:
        # 无限循环等待来自客户端连接的数据
        while True:
            # 使用loop.sock_recv异步接收数据
            data = await loop.sock_recv(connection, 1024)
            print(data)
            if not data:
                break

            # 一旦得到数据，将其发送回该客户端
            await loop.sock_sendall(connection, data)
    except ConnectionResetError:
        print("客户端断开连接")
    finally:
        connection.close()


echo_tasks = []


async def connection_listener(server_socket: socket, loop: AbstractEventLoop) -> None:
    while True:
        connection, address = await loop.sock_accept(server_socket)
        connection.setblocking(False)
        print("获取到连接")
        # 每当获取连接时，创建一个回显任务来监听客户端数据
        echo_task = asyncio.create_task(echo(connection, loop))
        echo_tasks.append(echo_task)


async def close_echo_tasks(echo_tasks: List[asyncio.Task]):
    waiters = [asyncio.wait_for(task, 2) for task in echo_tasks]
    for task in waiters:
        try:
            await task
        except asyncio.exceptions.TimeoutError:
            pass


class GracefulExit(SystemExit):
    pass


def shutdown():
    raise GracefulExit()


async def main():
    # 创建一个TCP服务器套接字
    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    #  将SO_REUSEADDR 设置为1，表示重复使用此端口，避免出现address already In use错误。
    server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    # 将套接字的地址设置为127.0.0.1:8000
    address = ("127.0.0.1", 8000)
    server_socket.setblocking(False)
    server_socket.bind(address)
    # 开启监听传入的连接
    server_socket.listen()

    for signame in {"SIGINT", "SIGTERM"}:
        loop.add_signal_handler(getattr(signal, signame), shutdown)
    # 启动协程以侦听连接
    await connection_listener(server_socket, asyncio.get_event_loop())


loop = asyncio.get_event_loop()

try:
    loop.run_until_complete(main())
except GracefulExit:
    loop.run_until_complete(close_echo_tasks(echo_tasks))
finally:
    loop.close()
